#include "lexer.h"
#include <cctype>
#include <stdexcept>

Lexer::Lexer(const std::string& source) 
    : source(source), position(0), line(1), column(1) {
    initializeKeywords();
}

void Lexer::initializeKeywords() {
    keywords = {
        {"int", TokenType::INT},
        {"char", TokenType::CHAR},
        {"float", TokenType::FLOAT},
        {"double", TokenType::DOUBLE},
        {"void", TokenType::VOID},
        {"if", TokenType::IF},
        {"else", TokenType::ELSE},
        {"while", TokenType::WHILE},
        {"for", TokenType::FOR},
        {"return", TokenType::RETURN},
        {"break", TokenType::BREAK},
        {"continue", TokenType::CONTINUE},
        {"struct", TokenType::STRUCT},
        {"typedef", TokenType::TYPEDEF},
        {"const", TokenType::CONST},
        {"static", TokenType::STATIC},
        {"extern", TokenType::EXTERN}
    };
}

char Lexer::currentChar() {
    if (position >= source.length()) {
        return '\0';
    }
    return source[position];
}

char Lexer::peekChar() {
    if (position + 1 >= source.length()) {
        return '\0';
    }
    return source[position + 1];
}

void Lexer::advance() {
    if (position < source.length()) {
        if (source[position] == '\n') {
            line++;
            column = 1;
        } else {
            column++;
        }
        position++;
    }
}

void Lexer::skipWhitespace() {
    while (std::isspace(currentChar()) && currentChar() != '\0') {
        advance();
    }
}

void Lexer::skipComment() {
    if (currentChar() == '/' && peekChar() == '/') {
        // Single line comment
        while (currentChar() != '\n' && currentChar() != '\0') {
            advance();
        }
    } else if (currentChar() == '/' && peekChar() == '*') {
        // Multi-line comment
        advance(); // skip /
        advance(); // skip *
        while (currentChar() != '\0') {
            if (currentChar() == '*' && peekChar() == '/') {
                advance(); // skip *
                advance(); // skip /
                break;
            }
            advance();
        }
    }
}

Token Lexer::scanNumber() {
    std::string value;
    int startLine = line;
    int startColumn = column;
    
    // Handle integer and float literals
    while (std::isdigit(currentChar()) || currentChar() == '.') {
        value += currentChar();
        advance();
    }
    
    // Determine if it's a float or integer
    if (value.find('.') != std::string::npos) {
        return Token(TokenType::FLOAT_LITERAL, value, startLine, startColumn);
    } else {
        return Token(TokenType::INTEGER_LITERAL, value, startLine, startColumn);
    }
}

Token Lexer::scanString() {
    std::string value;
    int startLine = line;
    int startColumn = column;
    
    advance(); // skip opening quote
    while (currentChar() != '"' && currentChar() != '\0') {
        if (currentChar() == '\\') {
            advance(); // skip backslash
            switch (currentChar()) {
                case 'n': value += '\n'; break;
                case 't': value += '\t'; break;
                case 'r': value += '\r'; break;
                case '\\': value += '\\'; break;
                case '"': value += '"'; break;
                default: value += currentChar(); break;
            }
        } else {
            value += currentChar();
        }
        advance();
    }
    advance(); // skip closing quote
    
    return Token(TokenType::STRING_LITERAL, value, startLine, startColumn);
}

Token Lexer::scanChar() {
    std::string value;
    int startLine = line;
    int startColumn = column;
    
    advance(); // skip opening quote
    if (currentChar() == '\\') {
        advance(); // skip backslash
        switch (currentChar()) {
            case 'n': value = "\n"; break;
            case 't': value = "\t"; break;
            case 'r': value = "\r"; break;
            case '\\': value = "\\"; break;
            case '\'': value = "'"; break;
            default: value = currentChar(); break;
        }
    } else {
        value = currentChar();
    }
    advance(); // skip character
    advance(); // skip closing quote
    
    return Token(TokenType::CHAR_LITERAL, value, startLine, startColumn);
}

Token Lexer::scanIdentifier() {
    std::string value;
    int startLine = line;
    int startColumn = column;
    
    while (std::isalnum(currentChar()) || currentChar() == '_') {
        value += currentChar();
        advance();
    }
    
    // Check if it's a keyword
    auto it = keywords.find(value);
    if (it != keywords.end()) {
        return Token(it->second, value, startLine, startColumn);
    }
    
    return Token(TokenType::IDENTIFIER, value, startLine, startColumn);
}

Token Lexer::scanOperator() {
    int startLine = line;
    int startColumn = column;
    
    char first = currentChar();
    advance();
    
    switch (first) {
        case '+':
            if (currentChar() == '+') {
                advance();
                return Token(TokenType::INCREMENT, "++", startLine, startColumn);
            }
            return Token(TokenType::PLUS, "+", startLine, startColumn);
            
        case '-':
            if (currentChar() == '-') {
                advance();
                return Token(TokenType::DECREMENT, "--", startLine, startColumn);
            } else if (currentChar() == '>') {
                advance();
                return Token(TokenType::ARROW, "->", startLine, startColumn);
            }
            return Token(TokenType::MINUS, "-", startLine, startColumn);
            
        case '*':
            return Token(TokenType::MULTIPLY, "*", startLine, startColumn);
            
        case '/':
            return Token(TokenType::DIVIDE, "/", startLine, startColumn);
            
        case '%':
            return Token(TokenType::MODULO, "%", startLine, startColumn);
            
        case '=':
            if (currentChar() == '=') {
                advance();
                return Token(TokenType::EQUAL, "==", startLine, startColumn);
            }
            return Token(TokenType::ASSIGN, "=", startLine, startColumn);
            
        case '!':
            if (currentChar() == '=') {
                advance();
                return Token(TokenType::NOT_EQUAL, "!=", startLine, startColumn);
            }
            return Token(TokenType::NOT, "!", startLine, startColumn);
            
        case '<':
            if (currentChar() == '=') {
                advance();
                return Token(TokenType::LESS_EQUAL, "<=", startLine, startColumn);
            }
            return Token(TokenType::LESS, "<", startLine, startColumn);
            
        case '>':
            if (currentChar() == '=') {
                advance();
                return Token(TokenType::GREATER_EQUAL, ">=", startLine, startColumn);
            }
            return Token(TokenType::GREATER, ">", startLine, startColumn);
            
        case '&':
            if (currentChar() == '&') {
                advance();
                return Token(TokenType::AND, "&&", startLine, startColumn);
            }
            return Token(TokenType::UNKNOWN, "&", startLine, startColumn);
            
        case '|':
            if (currentChar() == '|') {
                advance();
                return Token(TokenType::OR, "||", startLine, startColumn);
            }
            return Token(TokenType::UNKNOWN, "|", startLine, startColumn);
            
        case ';':
            return Token(TokenType::SEMICOLON, ";", startLine, startColumn);
            
        case ',':
            return Token(TokenType::COMMA, ",", startLine, startColumn);
            
        case '.':
            return Token(TokenType::DOT, ".", startLine, startColumn);
            
        case '(':
            return Token(TokenType::LEFT_PAREN, "(", startLine, startColumn);
            
        case ')':
            return Token(TokenType::RIGHT_PAREN, ")", startLine, startColumn);
            
        case '[':
            return Token(TokenType::LEFT_BRACKET, "[", startLine, startColumn);
            
        case ']':
            return Token(TokenType::RIGHT_BRACKET, "]", startLine, startColumn);
            
        case '{':
            return Token(TokenType::LEFT_BRACE, "{", startLine, startColumn);
            
        case '}':
            return Token(TokenType::RIGHT_BRACE, "}", startLine, startColumn);
            
        default:
            return Token(TokenType::UNKNOWN, std::string(1, first), startLine, startColumn);
    }
}

Token Lexer::nextToken() {
    while (currentChar() != '\0') {
        skipWhitespace();
        
        if (currentChar() == '\0') {
            break;
        }
        
        // Handle comments
        if (currentChar() == '/' && (peekChar() == '/' || peekChar() == '*')) {
            skipComment();
            continue;
        }
        
        if (std::isdigit(currentChar())) {
            return scanNumber();
        }
        
        if (currentChar() == '"') {
            return scanString();
        }
        
        if (currentChar() == '\'') {
            return scanChar();
        }
        
        if (std::isalpha(currentChar()) || currentChar() == '_') {
            return scanIdentifier();
        }
        
        return scanOperator();
    }
    
    return Token(TokenType::END_OF_FILE, "", line, column);
}

std::vector<Token> Lexer::tokenize() {
    std::vector<Token> tokens;
    Token token = nextToken();
    
    while (token.type != TokenType::END_OF_FILE) {
        tokens.push_back(token);
        token = nextToken();
    }
    
    tokens.push_back(token); // Add EOF token
    return tokens;
}
